import cv2
import numpy as np
import os
import csv
from datetime import datetime
from tkinter import Tk, Frame, Button, Label, Text, PhotoImage, Toplevel, WORD, Entry
from PIL import Image, ImageTk
import threading

# ---------------------- 系统常量 ----------------------
CASCADE_PATH = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
DATA_FOLDER = 'face_attendance_data'
MODEL_FILE = os.path.join(DATA_FOLDER, 'face_model.yml')
NAMES_FILE = os.path.join(DATA_FOLDER, 'names.csv')
ATTENDANCE_FILE = 'attendance_records.csv'
FACE_SIZE = (100, 100)  # 人脸样本尺寸
CAMERA_WIDTH = 640  # 降低摄像头分辨率以提升性能
CAMERA_HEIGHT = 480

# ---------------------- GUI 颜色配置 ----------------------
BG_COLOR = '#F0F0F0'
BUTTON_COLOR = '#4CAF50'
TEXT_COLOR = '#333333'
WARNING_COLOR = '#FF4444'
SUCCESS_COLOR = '#00C851'


class FaceAttendanceGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("人脸识别考勤系统")
        self.root.geometry("800x600")
        self.root.configure(bg=BG_COLOR)

        # 初始化组件
        self.face_cascade = None
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.known_names = []
        self.is_running = False  # 线程运行标志
        self.camera_thread = None  # 摄像头线程
        self.cap = None  # 摄像头对象

        self.load_data()
        self.create_widgets()

    def load_cascade(self):
        """延迟加载人脸检测器，提升启动速度"""
        if self.face_cascade is None:
            cascade = cv2.CascadeClassifier()
            if not cascade.load(CASCADE_PATH):
                self.show_warning("错误：无法加载人脸检测器文件")
                return None
            self.face_cascade = cascade
        return self.face_cascade

    def load_data(self):
        """加载已有数据"""
        try:
            if os.path.exists(MODEL_FILE):
                self.recognizer.read(MODEL_FILE)
            if os.path.exists(NAMES_FILE):
                with open(NAMES_FILE, 'r', encoding='utf-8') as f:
                    self.known_names = [line.strip() for line in f if line.strip()]
        except Exception as e:
            self.show_warning(f"警告：数据加载失败 - {str(e)}")

    def create_widgets(self):
        """创建GUI组件"""
        # 左侧功能区
        left_frame = Frame(self.root, bg=BG_COLOR, padx=20, pady=20)
        left_frame.pack(side="left", fill="y")

        Label(left_frame, text="人脸识别考勤系统", font=('Arial', 18, 'bold'), bg=BG_COLOR, fg=TEXT_COLOR).pack(pady=10)

        self.register_btn = Button(left_frame, text="员工注册", command=self.register_employee,
                                   bg=BUTTON_COLOR, fg='white', font=('Arial', 12), width=15)
        self.register_btn.pack(pady=10)

        self.attend_btn = Button(left_frame, text="开始考勤", command=self.start_attendance,
                                 bg=BUTTON_COLOR, fg='white', font=('Arial', 12), width=15)
        self.attend_btn.pack(pady=10)

        self.view_btn = Button(left_frame, text="查看考勤记录", command=self.view_records,
                               bg=BUTTON_COLOR, fg='white', font=('Arial', 12), width=15)
        self.view_btn.pack(pady=10)

        # 右侧视频显示区
        self.video_frame = Frame(self.root, bg=BG_COLOR, padx=20, pady=20)
        self.video_frame.pack(side="right", fill="both", expand=True)

        self.canvas = Label(self.video_frame, bg='black')
        self.canvas.pack(pady=20)

        # 状态提示栏
        self.status_label = Label(self.root, text="系统就绪", bg=BG_COLOR, fg=TEXT_COLOR, font=('Arial', 10))
        self.status_label.pack(pady=10)

    def show_warning(self, message):
        """显示警告信息"""
        self.status_label.config(text=message, fg=WARNING_COLOR)
        self.root.after(3000, lambda: self.status_label.config(text="系统就绪"))

    def show_success(self, message):
        """显示成功信息"""
        self.status_label.config(text=message, fg=SUCCESS_COLOR)
        self.root.after(3000, lambda: self.status_label.config(text="系统就绪"))

    def register_employee(self):
        """员工注册窗口"""
        reg_window = Toplevel(self.root)
        reg_window.title("员工注册")
        reg_window.geometry("400x300")
        reg_window.configure(bg=BG_COLOR)

        Label(reg_window, text="输入员工姓名:", font=('Arial', 12), bg=BG_COLOR, fg=TEXT_COLOR).pack(pady=10)
        name_entry = Entry(reg_window, font=('Arial', 12))
        name_entry.pack(pady=5)

        def start_registration():
            name = name_entry.get().strip()
            if not name:
                self.show_warning("姓名不能为空")
                return
            if name in self.known_names:
                self.show_warning("该员工已注册")
                return

            # 优化：使用独立线程处理摄像头，避免界面卡顿
            self.cap = cv2.VideoCapture(0)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)
            if not self.cap.isOpened():
                self.show_warning("无法打开摄像头")
                return

            samples = []
            reg_window.destroy()

            # 采集样本使用独立线程
            self.is_running = True
            self.camera_thread = threading.Thread(target=self.capture_samples, args=(name, self.cap, samples))
            self.camera_thread.daemon = True
            self.camera_thread.start()

        Button(reg_window, text="开始注册", command=start_registration,
               bg=BUTTON_COLOR, fg='white', font=('Arial', 12), width=10).pack(pady=10)

    def capture_samples(self, name, cap, samples):
        """使用独立线程采集人脸样本"""
        sample_count = 0
        while self.is_running and sample_count < 15:
            ret, frame = cap.read()
            if not ret:
                break

            # 优化：降低人脸检测频率，每3帧检测一次
            if cv2.getTickCount() % 3 == 0:
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                faces = self.load_cascade().detectMultiScale(gray, 1.3, 5)

                for (x, y, w, h) in faces:
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

                    # 显示采集进度
                    cv2.putText(frame, f"样本: {sample_count}/15", (10, 30),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

                    # 按空格键采集样本
                    key = cv2.waitKey(1)
                    if key == 32:  # 空格键
                        sample = gray[y:y + h, x:x + w]
                        sample = cv2.resize(sample, FACE_SIZE)
                        samples.append(sample)
                        sample_count += 1

                        # 更新主界面状态
                        self.root.after(0, lambda s=sample_count: self.show_success(f"已采集 {s}/15"))

            # 更新视频显示（降低更新频率）
            if cv2.getTickCount() % (cv2.getTickFrequency() // 10) == 0:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame)
                img_tk = ImageTk.PhotoImage(image=img)

                # 在主线程中更新UI
                self.root.after(0, lambda img=img_tk: self.canvas.config(image=img))
                self.canvas.img_tk = img_tk

        # 采集完成，释放资源
        self.is_running = False
        cap.release()
        cv2.destroyAllWindows()

        # 训练模型
        if len(samples) >= 10:
            self.root.after(0, lambda: self.train_model(name, samples))
        else:
            self.root.after(0, lambda: self.show_warning("采集样本不足，至少需要10张"))

    def train_model(self, name, samples):
        """训练人脸识别模型"""
        try:
            self.recognizer.update(samples, np.array([len(self.known_names)] * len(samples)))
            self.known_names.append(name)
            self.save_data()
            self.show_success(f"员工 {name} 注册成功")
        except Exception as e:
            self.show_warning(f"注册失败 - {str(e)}")

    def save_data(self):
        """保存数据"""
        try:
            if not os.path.exists(DATA_FOLDER):
                os.makedirs(DATA_FOLDER)
            self.recognizer.write(MODEL_FILE)
            with open(NAMES_FILE, 'w', encoding='utf-8') as f:
                f.write('\n'.join(self.known_names))
        except Exception as e:
            self.show_warning(f"数据保存失败 - {str(e)}")

    def start_attendance(self):
        """使用独立线程开始考勤打卡"""
        if self.is_running:
            return

        if self.cap is None or not self.cap.isOpened():
            self.cap = cv2.VideoCapture(0)
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)
            if not self.cap.isOpened():
                self.show_warning("无法打开摄像头")
                return

        self.is_running = True
        self.camera_thread = threading.Thread(target=self.run_attendance)
        self.camera_thread.daemon = True
        self.camera_thread.start()

    def run_attendance(self):
        """考勤处理线程"""
        checked = set()
        last_detection_time = 0

        while self.is_running:
            ret, frame = self.cap.read()
            if not ret:
                break

            current_time = cv2.getTickCount()

            # 优化：降低人脸检测频率，每5帧检测一次
            if (current_time - last_detection_time) / cv2.getTickFrequency() > 0.2:
                last_detection_time = current_time

                # 简化图像处理，关闭人脸美化
                # frame = self.beautify_face(frame)
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                faces = self.load_cascade().detectMultiScale(gray, 1.3, 5)

                for (x, y, w, h) in faces:
                    face_roi = gray[y:y + h, x:x + w]
                    face_roi = cv2.resize(face_roi, FACE_SIZE)

                    try:
                        fid, confidence = self.recognizer.predict(face_roi)
                    except:
                        continue

                    name = self.known_names[fid] if confidence < 50 else 'Unknown'
                    color = (0, 255, 0) if name != 'Unknown' else (0, 0, 255)

                    # 绘制识别结果（简化绘制参数以提升性能）
                    cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                    cv2.putText(frame, f"{name} ({100 - int(confidence)}%)",
                                (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1)

                    # 记录考勤（每天每人只记录一次）
                    if name != 'Unknown':
                        today = datetime.now().strftime('%Y-%m-%d')
                        key = f"{name}_{today}"
                        if key not in checked:
                            # 在主线程中执行考勤记录
                            self.root.after(0, lambda n=name: self.record_attendance(n))
                            checked.add(key)

            # 降低视频显示频率，每秒10帧
            if current_time % (cv2.getTickFrequency() // 10) == 0:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame)
                img_tk = ImageTk.PhotoImage(image=img)

                # 在主线程中更新UI
                self.root.after(0, lambda img=img_tk: self.canvas.config(image=img))
                self.canvas.img_tk = img_tk

        # 释放资源
        self.is_running = False
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()

    def record_attendance(self, name):
        """记录考勤"""
        now = datetime.now()
        try:
            with open(ATTENDANCE_FILE, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                if os.path.getsize(ATTENDANCE_FILE) == 0:
                    writer.writerow(['姓名', '日期', '时间', '类型'])
                writer.writerow([
                    name,
                    now.strftime('%Y-%m-%d'),
                    now.strftime('%H:%M:%S'),
                    '上班' if now.hour < 12 else '下班'
                ])
            self.show_success(f"{name} 考勤记录成功")
        except Exception as e:
            self.show_warning(f"记录失败 - {str(e)}")

    def view_records(self):
        """查看考勤记录"""
        if not os.path.exists(ATTENDANCE_FILE):
            self.show_warning("暂无考勤记录")
            return

        records_window = Toplevel(self.root)
        records_window.title("考勤记录")
        records_window.geometry("600x400")
        records_window.configure(bg=BG_COLOR)

        text_widget = Text(records_window, wrap=WORD, bg='white', fg=TEXT_COLOR, font=('Arial', 10))
        text_widget.pack(fill='both', expand=True, padx=10, pady=10)

        try:
            with open(ATTENDANCE_FILE, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                if reader.fieldnames:
                    header = f"{reader.fieldnames[0]:<10} {reader.fieldnames[1]:<15} {reader.fieldnames[2]:<10} {reader.fieldnames[3]}\n"
                    text_widget.insert('1.0', header)
                    text_widget.insert('1.0', '-' * 50 + '\n')
                    for row in reader:
                        line = f"{row['姓名']:<10} {row['日期']:<15} {row['时间']:<10} {row['类型']}\n"
                        text_widget.insert('end', line)
        except Exception as e:
            text_widget.insert('end', f"错误：{str(e)}")

    def beautify_face(self, frame):
        """简化的人脸美化（可选功能）"""
        try:
            # 仅保留基础磨皮，降低参数以提升性能
            return cv2.bilateralFilter(frame, 5, 10, 10)
        except:
            return frame

    def on_close(self):
        """关闭程序时清理资源"""
        self.is_running = False
        if self.camera_thread and self.camera_thread.is_alive():
            self.camera_thread.join(timeout=1.0)
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()
        self.root.destroy()


if __name__ == "__main__":
    root = Tk()
    app = FaceAttendanceGUI(root)
    root.protocol("WM_DELETE_WINDOW", app.on_close)  # 确保关闭窗口时释放资源
    root.mainloop()